import numpy as np
"""
一种计算均值的算法
"""
class wavelet_llt:
    def __init__(self,timeperiod=30):
        self.a=2/(timeperiod+1)
        self.data=0
        self.predata=[]
        self.pret=[]
        self.timekey=None
    def oncalc(self,t,timekey=None):
        return self.ondata(t,timekey)
    def ondata(self,t,timekey=None):
        if timekey:
            if timekey==self.timekey:
                self.predata.pop()
                self.pret.pop()
            else:
                self.timekey=timekey
        if t is not None:
            if len(self.predata)<2:
                self.data=t
                self.predata.append(self.data)
                self.pret.append(t)
            else:
                self.data=(self.a-self.a**2/4)*t+self.a**2/2*self.pret[-1]-(self.a-3*self.a**2/4)*self.pret[-2]+2*(1-self.a)*self.predata[-1]-(1-self.a)**2*self.predata[-2]
                self.predata.append(self.data)
                self.pret.append(t)
                while len(self.predata) > 3:
                    self.predata.pop(0)
                    self.pret.pop(0)
        return self.data
class wavelet_llt_new:
    def __init__(self,timeperiod=30):
        self.period = timeperiod
        self.a=2/(timeperiod+1)
        self.data=0
        self.predata=[]
        self.pret=[]
    def ondata(self,t):
        if len(self.pret)<self.period + 1:
            self.data=t
            self.predata = [0,0]
            self.pret.append(t)
        else:
            self.pret.append(t)
            #print('sys',len(self.pret))
            # temp1 = np.mean(self.pret[2:2 + self.period])
            # temp2 = np.mean(self.pret[1:1+self.period])
            # temp3 = np.mean(self.pret[0:self.period])
            # temp_1 = self.pret[2:2 + self.period]
            # temp_2 = self.pret[1:1 + self.period]
            # temp_3 = self.pret[0:self.period]
            self.data=(self.a-self.a**2/4)*np.mean(self.pret[2:2 + self.period])\
                      +self.a**2/2*np.mean(self.pret[1:1+self.period])-\
                      (self.a-3*self.a**2/4)*np.mean(self.pret[0:self.period])+\
                      2*(1-self.a)*self.predata[-1]-(1-self.a)**2*self.predata[-2]
            self.pret.pop(0)
            self.predata.pop(0)
            self.predata.append(self.data)
        return self.data
if __name__ == '__main__':
    from jili.core import load, save
    from research.calcor.calcors_graph import graph_calcor
    import pandas as pd
    k1m = load(r"D:\data\future_k1m_tq\TA005\TA005_20200102.pkl")
    roc_stat0=wavelet_llt(timeperiod=10)
    roc_stat1 = wavelet_llt_new(timeperiod=10)
    cc = [{'calc_cmd': 'ta', 'cmd': 'EMA', 'out': ['ema30'], 'input': {'price': 'close'}, 'arg': {'timeperiod': 30},
           'batch': 30},
          {'calc_cmd': 'ta', 'cmd': 'EMA', 'out': ['ema10'], 'input': {'price': 'close'}, 'arg': {'timeperiod': 10},
           'batch': 10},
          {'calc_cmd': 'function', 'function': roc_stat0.ondata, 'out': ['llt0'], 'input': ["close"]},
          {'calc_cmd': 'function', 'function': roc_stat1.ondata, 'out': ['llt1'], 'input': ["close"]},
          ]
    c = graph_calcor(cc)
    # k1ms = k1m[:32]
    for t in k1m:
        t = c.onbar(t)
    k1md=pd.DataFrame(k1m)
    # print(t)
    # t1 = k1m[32]
    # t1["timekey"] = t["timekey"]
    # t1 = c.onbar(t1)
    # print(t1)
    print("end")